use logos::Logos;

use super::{
    element::ElementTable,
    lexer::{ChemToken, ExprToken},
};

struct ExprParser {
    tokens: Vec<ExprToken>,
    pos: usize,
}

impl ExprParser {
    fn new(input: &str) -> Self {
        let tokens = ExprToken::lexer(input)
            .filter_map(|token| token.ok())
            .collect();
        Self { tokens, pos: 0 }
    }

    fn peek(&mut self) -> Option<&ExprToken> {
        self.tokens.get(self.pos)
    }

    fn next(&mut self) -> Option<&ExprToken> {
        self.pos += 1;
        self.tokens.get(self.pos - 1)
    }
}

pub fn parse(input: &str) -> Result<f64, String> {
    let mut lexer = ExprParser::new(input);
    expr_bp(&mut lexer, 0)
}

fn expr_bp(lexer: &mut ExprParser, min_bp: u8) -> Result<f64, String> {
    let mut lhs = match lexer.next() {
        Some(ExprToken::Number(num)) => *num,
        Some(ExprToken::Lparent) => {
            let lhs = expr_bp(lexer, 0)?;
            if lexer.next() != Some(&ExprToken::Rparent) {
                return Err("expect )".to_string());
            }
            lhs
        }
        Some(ExprToken::Chemical(chemical)) => chem_parse(&chemical[1..&chemical.len() - 1])?,
        t => return Err(format!("bad token: {:?}", t)),
    };
    loop {
        let op = match lexer.peek() {
            None | Some(ExprToken::Rparent) => break, // EOF
            Some(ExprToken::Add) => '+',
            Some(ExprToken::Sub) => '-',
            Some(ExprToken::Mul) => '*',
            Some(ExprToken::Div) => '/',
            t => return Err(format!("bad token: {:?}", t)),
        };

        if let Some((l_bp, r_bp)) = infix_binding_power(op) {
            if l_bp < min_bp {
                break;
            }
            lexer.next();
            lhs = {
                let rhs = expr_bp(lexer, r_bp)?;
                match op {
                    '+' => lhs + rhs,
                    '-' => lhs - rhs,
                    '*' => lhs * rhs,
                    '/' => lhs / rhs,
                    t => panic!("bad token: {:?}", t),
                }
            }
        }
    }
    Ok(lhs)
}

fn infix_binding_power(op: char) -> Option<(u8, u8)> {
    let res = match op {
        '+' | '-' => (5, 6),
        '*' | '/' => (7, 8),
        _ => return None,
    };
    Some(res)
}

struct ChemParser {
    tokens: Vec<ChemToken>,
    pos: usize,
}

impl ChemParser {
    fn new(input: &str) -> Self {
        let input = input.trim();
        let tokens: Vec<ChemToken> = ChemToken::lexer(input)
            .filter_map(|token| token.ok())
            .collect();

        // delet unnessary space
        let tokens = tokens
            .iter()
            .enumerate()
            .filter_map(|(idx, iter)| match iter {
                ChemToken::Space => {
                    if idx > 0 && idx < tokens.len() - 1 {
                        let prev = &tokens[idx - 1];
                        let next = &tokens[idx + 1];
                        if prev == &ChemToken::Lparent || next == &ChemToken::Rparent {
                            None
                        } else {
                            Some(*iter)
                        }
                    } else {
                        None
                    }
                }
                _ => Some(*iter),
            })
            .collect();

        Self { tokens, pos: 0 }
    }

    fn peek(&mut self) -> Option<&ChemToken> {
        self.tokens.get(self.pos)
    }

    fn next(&mut self) -> Option<&ChemToken> {
        self.pos += 1;
        self.tokens.get(self.pos - 1)
    }
}

pub fn chem_parse(chemical: &str) -> Result<f64, String> {
    let mut lexer = ChemParser::new(chemical);
    parse_total(&mut lexer)
}

fn parse_total(lexer: &mut ChemParser) -> Result<f64, String> {
    let mut res = parse_prefix_group(lexer)?;
    match lexer.peek() {
        Some(ChemToken::Space) => {
            lexer.next();
            res += parse_total(lexer)?;
        }
        Some(ChemToken::Rparent) | None => (),
        t => return Err(format!("bad token: {:?}", t)),
    }
    Ok(res)
}

fn parse_prefix_group(lexer: &mut ChemParser) -> Result<f64, String> {
    let mut res = match lexer.peek() {
        Some(ChemToken::Number(num)) => {
            let res = *num;
            lexer.next();
            res
        }
        _ => 1.0,
    };
    res *= parse_group(lexer)?;
    Ok(res)
}

fn parse_group(lexer: &mut ChemParser) -> Result<f64, String> {
    let mut res = parse_atom(lexer)?;
    if let Some(ChemToken::Number(num)) = lexer.peek() {
        res *= num;
        lexer.next();
    }
    match lexer.peek() {
        Some(ChemToken::Space) | Some(ChemToken::Rparent) | None => (),
        _ => res += parse_group(lexer)?,
    }
    Ok(res)
}
fn parse_atom(lexer: &mut ChemParser) -> Result<f64, String> {
    match lexer.peek() {
        Some(ChemToken::Lparent) => {
            lexer.next();
            let mut res = parse_total(lexer)?;
            assert_eq!(lexer.next(), Some(&ChemToken::Rparent));
            if let Some(ChemToken::Number(num)) = lexer.peek() {
                res *= num;
                lexer.next();
            }
            Ok(res)
        }
        Some(ChemToken::Elem(elem)) => {
            let ans = parse_element(*elem)?;
            lexer.next();
            Ok(ans)
        }
        _ => Ok(0.0),
    }
}
fn parse_element(elem: (char, char)) -> Result<f64, String> {
    ElementTable::new().get_mw(elem).ok_or(format!("{}{} is not an element", elem.0, elem.1))
}
